#include "stun.h"

#include <rtc_base/byte_order.h>
#include <rtc_base/crc32.h>
#include <rtc_base/message_digest.h>
#include <rtc_base/logging.h>

namespace xrtc {

const char EMPTY_TRANSACION_ID[] = "000000000000"; // 占用12个字节
const size_t STUN_FINGERPRINT_XOR_VALUE = 0x5354554e;

// STUN错误描述信息
const char STUN_ERROR_REASON_BAD_REQUEST[] = "Bad request";
const char STUN_ERROR_REASON_UNATHORIZED[] = "Unathorized";
const char STUN_ERROR_REASON_SERVER_ERROR[] = "Server error";


std::string stun_method_to_string(int type) {
    switch (type) {
        case STUN_BINDING_REQUEST:
            return "BINDING REQUEST";
        case STUN_BINDING_RESPONSE:
            return "BINDING RESPONSE";
        default:
            return "Unknown<" + std::to_string(type) + ">";
    }
}

StunAttribute::StunAttribute(uint16_t type, uint16_t length) : m_type(type), m_length(length) {

}

StunByteStringAttribute::StunByteStringAttribute(uint16_t type, uint16_t length) : StunAttribute(type, length) {

}

StunByteStringAttribute::StunByteStringAttribute(uint16_t type, const std::string &str) : StunAttribute(type, 0) {
    copy_bytes(str.c_str(), str.size());
}

StunByteStringAttribute::~StunByteStringAttribute() {
    if (m_bytes) {
        delete[] m_bytes;
        m_bytes = nullptr;
    }
}

void StunByteStringAttribute::copy_bytes(const char *bytes, size_t length) {
    char* new_bytes = new char[length];
    memcpy(new_bytes, bytes, length);
    _set_bytes(new_bytes);
    set_length(length);
}

void StunByteStringAttribute::_set_bytes(char* bytes) {
    if (m_bytes) {
        delete[] m_bytes;
    }
    m_bytes = bytes;
}

bool StunByteStringAttribute::read(rtc::ByteBufferReader *buf)
{
    m_bytes = new char[getLength()];
    if (!buf->ReadBytes(m_bytes, getLength())) {
        return false;
    }
    
    consume_padding(buf);
    
    return true;
}

bool StunByteStringAttribute::write(rtc::ByteBufferWriter *buf) {
    buf->WriteBytes(m_bytes, getLength());
    write_padding(buf);
    return true;
}

// 工厂方法：根据传入的类型和长度创建STUN属性
StunAttribute *StunAttribute::createStunAttribute(StunAttributeValueType value_type, uint16_t type, uint16_t length, void *owner)
{
    // TODO
    switch (value_type) {
        case STUN_VALUE_BYTE_STRING:
            return new StunByteStringAttribute(type, length);
        case STUN_VALUE_UINT32:
            return new StunUInt32Attribute(type);
        default:
            return nullptr;
    }
} 

// 处理4字节偏移对齐的情况
void StunAttribute::consume_padding(rtc::ByteBufferReader *buf) {
    int remain = getLength() % 4;
    if (remain > 0) {
        buf->Consume(4 - remain);
    }
}

// 写入4字节偏移对齐的情况
void StunAttribute::write_padding(rtc::ByteBufferWriter *buf) {
    int remain = getLength() % 4;
    if (remain > 0) {
        
        char zeroes[4] = {0};
        // 填充4字节偏移对齐的情况
        buf->WriteBytes(zeroes, 4 - remain);
    }
}

StunMessage::StunMessage() : m_message_type(0), m_message_length(0), m_transaction_id(EMPTY_TRANSACION_ID) {

}


// 验证Fingerprint属性
bool StunMessage::validate_fingerprint(const char *data, size_t size) {
    // ========== 第一步：基本长度和格式检查 ==========
    // STUN消息必须是4字节对齐的，且包含fingerprint属性
    size_t fingerprint_attr_size = k_stun_attribute_header_size +  // 属性头部(4字节)
        StunUInt32Attribute::SIZE;  // fingerprint属性值(4字节)
    
    // 检查消息长度：必须是4字节对齐，且至少包含STUN头部+fingerprint属性
    if (size % 4 != 0 || size < k_stun_header_size + fingerprint_attr_size) {
        return false;
    }
    
    // ========== 第二步：验证STUN Magic Cookie ==========
    // STUN消息的8-11字节必须是固定的Magic Cookie (0x2112A442)
    // 这个值用于区分STUN消息和其他协议消息
    const char* magic_cookie = data + k_stun_transaction_id_offset -  // 获取Magic Cookie起始偏移
        k_stun_magic_cookie_length;  // 向前偏移4字节到第4字节位置
    if (rtc::GetBE32(magic_cookie) != k_stun_magic_cookie) {
        return false;  // Magic Cookie不匹配，不是有效的STUN消息
    }
    
    // ========== 第三步：验证Fingerprint属性位置和类型 ==========
    // fingerprint属性必须是STUN消息的最后一个属性
    const char* fingerprint_attr_data = data + size - fingerprint_attr_size;
    
    // 检查属性类型：必须是FINGERPRINT (0x8028)
    if (rtc::GetBE16(fingerprint_attr_data) != STUN_ATTR_FINGERPRINT ||
        // 检查属性长度：必须是4字节
        rtc::GetBE16(fingerprint_attr_data + sizeof(uint16_t)) != StunUInt32Attribute::SIZE) 
    {
        return false;  // fingerprint属性类型或长度不正确
    }
    
    // ========== 第四步：验证Fingerprint校验值 ==========
    // 获取fingerprint属性的实际值（4字节）
    uint32_t fingerprint = rtc::GetBE32(fingerprint_attr_data + 
            k_stun_attribute_header_size);

    // fingerprint校验算法：
    // 1. 计算STUN消息（除fingerprint属性外）的CRC32校验和
    // 2. 将CRC32与XOR值(0x5354554e)进行异或运算
    // 3. 结果应该等于fingerprint属性的值
    return (fingerprint ^ STUN_FINGERPRINT_XOR_VALUE) ==
        rtc::ComputeCrc32(data, size - fingerprint_attr_size);
}


// 添加Fingerprint属性
bool StunMessage::add_fingerprint() {
    auto fingerprint_attr_ptr = std::make_unique<StunUInt32Attribute>(STUN_ATTR_FINGERPRINT, 0);
    auto fingerprint_attr = fingerprint_attr_ptr.get();
    add_attribute(std::move(fingerprint_attr_ptr));

    rtc::ByteBufferWriter buf;
    if (!write(&buf)) {
        return false;
    }

    size_t msg_len_for_crc32 = buf.Length() - k_stun_attribute_header_size - fingerprint_attr->getLength();
    uint32_t crc32 = rtc::ComputeCrc32(buf.Data(), msg_len_for_crc32);
    fingerprint_attr->set_value(crc32 ^ STUN_FINGERPRINT_XOR_VALUE);

    return true;
}

// 验证消息完整性的状态
StunMessage::MessageIntegrityStatus StunMessage::validate_message_integrity(const std::string &pwd) {
    m_pwd = pwd;
    if (get_byte_string_attribute(STUN_ATTR_MESSAGE_INTEGRITY)) {
        if (_validate_message_integrity_of_type(STUN_ATTR_MESSAGE_INTEGRITY, 
            k_stun_message_integrity_size, m_buffer.c_str(), m_buffer.size(), pwd)) {
            m_message_integrity_status = MessageIntegrityStatus::k_integrity_ok;
        } else {
            return MessageIntegrityStatus::k_integrity_bad;
        }
    } else {
        m_message_integrity_status = MessageIntegrityStatus::k_no_integrity;
    }

    return m_message_integrity_status;
}

// 添加消息完整性属性
bool StunMessage::add_message_integrity(const std::string& password) {
    return _add_message_integrity_of_type(STUN_ATTR_MESSAGE_INTEGRITY,
        k_stun_message_integrity_size, password.c_str(),
        password.size());
}

// 添加消息完整性属性
bool StunMessage::_add_message_integrity_of_type(uint16_t attr_type, uint16_t attr_size, 
                                        const char* pwd, size_t pwd_length) {
    auto mi_attr_ptr = std::make_unique<StunByteStringAttribute>(attr_type, std::string(attr_size, '0'));
    auto mi_attr = mi_attr_ptr.get();
    add_attribute(std::move(mi_attr_ptr));

    rtc::ByteBufferWriter buf;
    if (!write(&buf)) {
        return false;
    }

    size_t msg_len_for_hmac = buf.Length() - k_stun_attribute_header_size - mi_attr->getLength();
    char hmac[k_stun_message_integrity_size];
    size_t ret = rtc::ComputeHmac(rtc::DIGEST_SHA_1, pwd, pwd_length, buf.Data(), msg_len_for_hmac, hmac, sizeof(hmac));
    if (ret != sizeof(hmac)) {
        RTC_LOG(LS_WARNING) << "compute hmac error";
        return false;
    }

    mi_attr->copy_bytes(hmac, k_stun_message_integrity_size);
    m_pwd.assign(pwd, pwd_length);
    m_message_integrity_status = MessageIntegrityStatus::k_integrity_ok;

    return true;
}

// 添加STUN属性
void StunMessage::add_attribute(std::unique_ptr<StunAttribute> attr) {
    size_t attr_len = attr->getLength();
    if (attr_len % 4 != 0) {
        // 不满足4字节对齐，需要补齐
        attr_len += (4 - (attr_len % 4));
    }

    m_message_length += (attr_len + k_stun_attribute_header_size);
    m_attributes.push_back(std::move(attr));
}

// 获取STUN属性值类型
StunAttributeValueType StunMessage::get_attribute_value_type(int type) {
    switch (type) {
        // // 用户名属性
        case STUN_ATTR_USERNAME:
            return STUN_VALUE_BYTE_STRING;
        // 消息完整性属性
        case STUN_ATTR_MESSAGE_INTEGRITY:
            return STUN_VALUE_BYTE_STRING;
        case STUN_ATTR_PRIORITY:
            return STUN_VALUE_UINT32;
        // 未知类型
        default:
            return STUN_VALUE_UNKNOWN;
    }
}

// 读取STUN消息
bool StunMessage::read_from(rtc::ByteBufferReader *buf)
{
    if (!buf) {
        return false;
    }

    // 保存STUN消息内容的缓冲区，用于验证消息完整性
    m_buffer.assign(buf->Data(), buf->Data() + buf->Length());
    
    if (!buf->ReadUInt16(&m_message_type)) {
        return false;
    }
    
    // rtp/rtcp 10(2) 过滤rtp/rtcp
    if (m_message_type & 0x0800) {
        return false;
    }
    
    if (!buf->ReadUInt16(&m_message_length)) {
        return false;
    }
    
    std::string magic_cookie;
    if (!buf->ReadString(&magic_cookie, k_stun_magic_cookie_length)) {
        return false;
    }
    
    std::string transaction_id;
    if (!buf->ReadString(&transaction_id, k_stun_transaction_id_length)) {
        return false;
    }
    
    uint32_t magic_cookie_int;
    memcpy(&magic_cookie_int, magic_cookie.data(), sizeof(magic_cookie_int));
    if (rtc::NetworkToHost32(magic_cookie_int) != k_stun_magic_cookie) {
        transaction_id.insert(0, magic_cookie);
    }
    
    m_transaction_id = transaction_id;
    
    if (buf->Length() != m_message_length) {
        return false;
    }
    
    m_attributes.resize(0);

    while (buf->Length() > 0) {
        uint16_t attr_type;
        uint16_t attr_length;
        if (!buf->ReadUInt16(&attr_type)) {
            return false;
        }

        if (!buf->ReadUInt16(&attr_length)) {
            return false;
        }
        
        // 使用工厂模式创建属性对象
        std::unique_ptr<StunAttribute> attr(_create_attribute(attr_type, attr_length));
        if (!attr) {
            if (attr_length % 4 != 0) {
                attr_length += (4 - (attr_length % 4));
            }

            if (!buf->Consume(attr_length)) {
                return false;
            }
        } else {
            if (!attr->read(buf)) {
                return false;
            }

            m_attributes.push_back(std::move(attr));
        }
    }

    return true;
}

bool StunMessage::write(rtc::ByteBufferWriter *buf) const
{
    if (!buf) {
        return false;
    }

    // 写入STUN消息头部
    buf->WriteUInt16(m_message_type);
    buf->WriteUInt16(m_message_length);
    buf->WriteUInt32(k_stun_magic_cookie);
    buf->WriteString(m_transaction_id);
    
    // 遍历STUN属性，写入STUN属性
    for (const auto& attr : m_attributes) {
        buf->WriteUInt16(attr->getType());
        buf->WriteUInt16(attr->getLength());
        if (!attr->write(buf)) {
            return false;
        }
    }

    return true;
}

const StunUInt32Attribute *StunMessage::get_uint32(uint16_t type) {
    return static_cast<const StunUInt32Attribute*>(_get_attribute(type));
}

// 获取字符串类型STUN属性
const StunByteStringAttribute *StunMessage::get_byte_string_attribute(uint16_t type) {
    return static_cast<const StunByteStringAttribute*>(_get_attribute(type));
}

// 根据传入的类型和长度调用工厂方法创建属性对象
StunAttribute* StunMessage::_create_attribute(uint16_t type, uint16_t length) {
    // 1. 根据属性类型确定值类型
    StunAttributeValueType value_type = get_attribute_value_type(type);
    if (value_type != STUN_VALUE_UNKNOWN) {
        // 2. 调用工厂方法创建属性对象
        return StunAttribute::createStunAttribute(value_type, type, length, this);
    }
    return nullptr;
}

// 获取指定Type的STUN属性
const StunAttribute *StunMessage::_get_attribute(uint16_t type) {
    for (const auto& attr : m_attributes) {
        if (attr->getType() == type) {
            return attr.get();
        }
    }
   
    return nullptr;
}

// 解析MIP属性，验证消息完整性的状态
/*
    0                   1                   2                   3
 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|      STUN Message Type        |         Message Length        | → 消息类型和长度
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                         Magic Cookie                          | → 固定值0x2112A442
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                                                               |
|                     Transaction ID (96 bits)                 | → 事务ID
|                                                               |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
|                                                               |
|                         Attributes                            | → 属性列表
|                          (variable)                          |
|                                                               |
+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
 */
bool StunMessage::_validate_message_integrity_of_type(uint16_t mi_attr_type, size_t mi_attr_size, 
                        const char *data, size_t size, const std::string &pwd) {
    // 1. 基本长度和格式检查(是否是4字节对齐的，且至少包含STUN头部+MI属性)
    if (size % 4 != 0 || size < k_stun_header_size) {
        return false;
    }
    
    // 2. 检查属性长度 (前2个字节是Message Type)
    uint16_t length = rtc::GetBE16(&data[2]); // STUN消息长度
    if (length + k_stun_header_size != size) {
        return false;
    }

    // 查找MI属性的位置
    size_t current_pos = k_stun_header_size;
    bool has_message_integrity_status = false;
    while (current_pos + k_stun_attribute_header_size <= size) {
        // 3. 检查STUN属性类型和长度
    /*
          0                   1                   2                   3
          0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
          +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
          |             Type              |          Length               |
          +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
          |                             Value                             |
          |                             ...                               |
          +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
    */
        uint16_t attr_type = rtc::GetBE16(&data[current_pos]); // STUN属性类型
        uint16_t attr_length = rtc::GetBE16(&data[current_pos + sizeof(attr_type)]); // STUN属性长度
        if (attr_type == mi_attr_type) {
            has_message_integrity_status = true;
            break;
        }

        // 4. 偏移属性头部和属性值，处理4字节偏移对齐的情况(可能存在填充的字节)
        current_pos += (k_stun_attribute_header_size + attr_length);
        if (attr_length % 4 != 0) {
            current_pos += (4 - (attr_length % 4));
        }
    }
    
    if (!has_message_integrity_status) {
        return false;
    }

    // 5. 验证MI属性值
    size_t mi_pos = current_pos; // MI属性起始位置
    std::unique_ptr<char[]> temp_data(new char[mi_pos]);
    // 5.1 复制MI属性值之前的STUN消息内容到临时数据
    memcpy(temp_data.get(), data, mi_pos); 
    // 判断收到的 STUN 消息中，在 MESSAGE-INTEGRITY (MI) 属性之后是否还存在其他属性（如 FINGERPRINT）
    // 从而决定是否需要调整消息长度字段以进行正确的 HMAC 计算
    if (size > current_pos + k_stun_attribute_header_size + mi_attr_size) {
        // 5.2 计算MI属性值之后的STUN消息内容长度
        size_t extra_pos = mi_pos + k_stun_attribute_header_size + mi_attr_size;
        // 5.3 // MI之后的数据长度
        size_t extra_size = size - extra_pos;
        // 5.4 计算参与 HMAC 计算的属性部分长度
        // 临时修正了STUN消息头中的长度字段，使接收方在验证消息完整性时，能够基于与发送方完全相同的原始数据来计算HMAC
        // 仅包含 STUN 头部和 MI 属性之前所有属性的“逻辑上的”消息长度
        size_t adjust_new_len = size - extra_size - k_stun_header_size;
        rtc::SetBE16(temp_data.get() + 2, adjust_new_len);
    }

    // 6. 计算HMAC MESSAGE-INTEGRITY 计算之后的摘要（Digest）是 20 个字节
    char hmac[k_stun_message_integrity_size];
    size_t ret = rtc::ComputeHmac(rtc::DIGEST_SHA_1, pwd.c_str(), pwd.length(), temp_data.get(), 
                                  current_pos, hmac, sizeof(hmac));
    
    if (ret != k_stun_message_integrity_size) {
        return false;
    }

    // 7. 比较HMAC
    return memcmp(data + mi_pos + k_stun_attribute_header_size, hmac, mi_attr_size) == 0;
}

bool StunUInt32Attribute::read(rtc::ByteBufferReader *buf) {
    if (getLength() != SIZE || !buf->ReadUInt32(&m_bits)) {
        return false;
    }
    return true;
}

bool StunUInt32Attribute::write(rtc::ByteBufferWriter *buf) {
    buf->WriteUInt32(m_bits);
    return true;
}

StunAddressAttribute::StunAddressAttribute(uint16_t type, const rtc::SocketAddress &addr)
    : StunAttribute(type, 0) {
    set_address(addr);
}

void StunAddressAttribute::set_address(const rtc::SocketAddress &addr) {
    m_address = addr;
    
    switch (get_family()) {
        case STUN_ADDRESS_IPV4:
            set_length(SIZE_IPV4);
            break;
        case STUN_ADDRESS_IPV6:
            set_length(SIZE_IPV6);
            break;
        default:
            set_length(SIZE_UNDEF);
            break;
    }
}

StunAddressFamily StunAddressAttribute::get_family()
{
    switch (m_address.family()) {
        case AF_INET:
            return STUN_ADDRESS_IPV4;
        case AF_INET6:
            return STUN_ADDRESS_IPV6;
        default:
            return STUN_ADDRESS_UNDEF;
    }
}

bool StunAddressAttribute::read(rtc::ByteBufferReader *buf)
{
    // TODO
    return false;
}

bool StunAddressAttribute::write(rtc::ByteBufferWriter *buf)
{
    StunAddressFamily stun_address_family = get_family();
    if (stun_address_family == STUN_ADDRESS_UNDEF) {
        RTC_LOG(LS_WARNING) << "write address attribute error: unknown family";
        return false;
    }

    // 写入MAPPED-ADDRESS属性
    buf->WriteUInt8(0);
    buf->WriteUInt8(stun_address_family);
    buf->WriteUInt16(m_address.port());

    switch (m_address.family()) {
        case AF_INET: {
            in_addr v4_addr = m_address.ipaddr().ipv4_address();
            buf->WriteBytes((const char*)(&v4_addr), sizeof(v4_addr));
            break;
        }
        case AF_INET6: {
            in6_addr v6_addr = m_address.ipaddr().ipv6_address();
            buf->WriteBytes((const char*)(&v6_addr), sizeof(v6_addr));
            break;
        }
        default:
            RTC_LOG(LS_WARNING) << "write address attribute error: unknown family";
            return false;
    }

    return true;
}

StunXorAddressAttribute::StunXorAddressAttribute(uint16_t type, const rtc::SocketAddress &addr)
    : StunAddressAttribute(type, addr) {
}

bool StunXorAddressAttribute::write(rtc::ByteBufferWriter *buf) {
    StunAddressFamily stun_address_family = get_family();
    if (stun_address_family == STUN_ADDRESS_UNDEF) {
        RTC_LOG(LS_WARNING) << "write address attribute error: unknown family";
        return false;
    }

    rtc::IPAddress xored_ip = _get_xored_ip();
    if (xored_ip.family() == AF_UNSPEC) {
        return false;
    }

    // 写入XOR-MAPPED-ADDRESS属性
    buf->WriteUInt8(0);
    buf->WriteUInt8(stun_address_family);
    buf->WriteUInt16(m_address.port() ^ (k_stun_magic_cookie >> 16));

    switch (m_address.family()) {
        case AF_INET: {
            in_addr v4_addr = xored_ip.ipv4_address();
            buf->WriteBytes((const char*)(&v4_addr), sizeof(v4_addr));
            break;
        }
            
        case AF_INET6: {
            in6_addr v6_addr = m_address.ipaddr().ipv6_address();
            buf->WriteBytes((const char*)(&v6_addr), sizeof(v6_addr));
            break;
        }
            
        default:
            RTC_LOG(LS_WARNING) << "write address attribute error: unknown family";
            return false;
    }

    return true;
}

// 获取XOR-IP地址
rtc::IPAddress StunXorAddressAttribute::_get_xored_ip() {
    rtc::IPAddress ip = m_address.ipaddr();
    switch (m_address.family()) {
        case AF_INET: {
            in_addr v4addr = ip.ipv4_address();
            v4addr.s_addr = (v4addr.s_addr ^ rtc::HostToNetwork32(k_stun_magic_cookie));
            return rtc::IPAddress(v4addr);
        }
        case AF_INET6:
            break;
        default:
            break;
    }

    return rtc::IPAddress();
}

} // namespace xrtc